#include <iostream>
#include <fstream>
#include <string>
#include <random>
#include <set>
#include <map>
#include <vector>
#include <boost/algorithm/string.hpp>
#include <mpi.h>
#include "lp_lib.h"

using namespace std;

// Message tags
// 101: incremental graph: nodes (0 to the others)
// 102: incremental graph: edges (0 to the others)
// 201: nodes and attaced edges (all to all)
// 202: state moved with nodes (all to all)
// 203: nodes and processors they belong to (all to 0)
// 301: low-level (movable) nodes (all to 0)
// 302: optimal numbers of nodes that should be moved (0 to all)
// 401: request positions of boundary nodes (all to 0)
// 402: positions of boundary nodes (0 to all)
// 501: state of a node (all to all)

default_random_engine gen;
poisson_distribution<int> rpois0(1);

bool timeit = true;
bool use_label = false;
bool verbose = false;

void outedges(map<int, set<int>> &graph, int rank) {
    int j = 0;
    for (const auto &x : graph)
    for (int i : x.second)
    if (graph.find(i) == graph.end()) {
        cout << rank << ": (" << x.first << ", " << i << ") ";
        j++;
    }
    cout << "\n";
    cout << "Processor " << rank << " has " << j << " out-edges\n";
}

void print(map<int, set<int>> &graph, int rank) {
    cout << "Processor " << rank << " has " << graph.size() << " nodes ";
    for (const auto &g : graph)    cout << g.first << ", ";
    cout << "\n";
}
void output(string filename, map<int, int> position) {
    if (verbose) cout << "Writing to file " << filename << "...\n";

    ofstream output;
    output.open(filename);
    map<int, int> count;

    for (const auto &vertex : position) {
        output << vertex.first << " " << vertex.second << "\n";
        
        if (count.find(vertex.second) == count.end())
            count[vertex.second] = 0;
        count[vertex.second]++;
    }
    
}

// Read initial graph
// In proc 0, construct global table and local boundary set
void readGraph(string filename, int rank, int size,
        map<int, set<int>> &graph,
        map<int, int> &position) {
    
    ifstream init_graph (filename);
    string line;
    if (init_graph.is_open()) {
        while (getline(init_graph, line)) {
            vector<string> strs;
            boost::split(strs, line, boost::is_space());
            int source = stoi(strs[0]);
            int proc = stoi(strs[1]) % size;
            if (proc == rank) {
                graph[source] = set<int>();
                for (int i=2; i<strs.size(); i++) {
                    int target = stoi(strs[i]);
                    graph[source].insert(target);
                }
            }
            if (rank == 0)
            position[source] = proc;
        }
    }
    
    init_graph.close();
    if (verbose) cout << "Graph has been loaded in processor " << rank << ".\n";
}

void addGraph(vector<int> &added, vector<int> *added_state,
        map<int, set<int>> &graph, map<int, int> &state,
        int rank) {
    
    int idx = 0;
    while (idx < added.size()) {
        int curr = added[idx++];
        state[curr] = 0; (added_state == NULL) ? rpois0(gen) : (*added_state)[curr];
        if (graph.find(curr) == graph.end())
            graph[curr] = set<int>();
        while (added[idx++] > -1);
    }
    
    idx = 0;
    while (idx < added.size()) {
        int curr = added[idx++];
        while (added[idx] > -1) {
            int target = added[idx++];
            graph[curr].insert(target);
            if (graph.find(target) != graph.end())
                graph[target].insert(curr);
        }
        idx++;
    }
}

void addEdges(vector<int> &edges, map<int, set<int>> &graph) {
    for (int i=0; i<edges.size(); i+=2)
        graph[edges[i]].insert(edges[i+1]);
}

// Read graph in processor 0, find their location and boundary
void readIncremental(string filename, map<int, set<int>> &graph,
        map<int, int> &boundary, map<int, int> &position,
        map<int, int> &state, int size) {
    
    map<int, set<int>> newgraph;
    vector<vector<int>> loc(size, vector<int>());
    vector<vector<int>> edges(size, vector<int>());
    
    map<int, int> labels;
    ifstream add_graph (filename);
    string line;
    
    if (add_graph.is_open())
        while (getline(add_graph, line)) {
            vector<string> strs;
            boost::split(strs, line, boost::is_space());
            int source = stoi(strs[0]);
            
            if (use_label)
                labels[source] = stoi(strs[1]) % size;
            
            newgraph[source] = set<int>();
            for (int i=2; i<strs.size(); i++)
                newgraph[source].insert(stoi(strs[i]));
 
        }
    
    add_graph.close();
    
 
    for (const auto &vertex : newgraph) {
        int nearest = 0;
        
        if (use_label)
            nearest = labels[vertex.first];
        else {
            int deg = 0;
            vector<int> count(size, 0);
            
            for (const auto target : vertex.second) {
                count[position[target]]++;
                deg++;
            }
            
            int num = rand() % deg;
            for (int i=0; i<size; i++) {
                num -= count[i];
                if (num < 0) {
                    nearest = i;
                    break;
                }
            }
        }
        
        loc[nearest].push_back(vertex.first);
        
        for (int target : vertex.second) 
            loc[nearest].push_back(target);
        
        
        loc[nearest].push_back(-1);
        position[vertex.first] = nearest;
        
    }
    
    for (const auto &vertex : newgraph) {
        int p = position[vertex.first];
        for (const auto target : vertex.second) {
            int q = position[target];
            if (p != q) {
                edges[q].push_back(target);
                edges[q].push_back(vertex.first);
            }
        }
    }
    
    
    addGraph(loc[0], NULL, graph, state, 0);
    addEdges(edges[0], graph);
    
    for (int i=1; i<size; i++) {
        MPI_Send(&loc[i][0], loc[i].size(), MPI_INT, i, 101, MPI_COMM_WORLD);
        MPI_Send(&edges[i][0], edges[i].size(), MPI_INT, i, 102, MPI_COMM_WORLD);
    }
    
    if (verbose) cout << "A subgraph has been loaded in processor 0.\n";
    
}

void updateLayer(map<int, map<int, set<int>>> &layer,
        map<int, set<int>> &graph, map<int, int> &boundary,
        vector<int> *bd_size, int rank, int level) {
    
    for (int k = 0; k < level; k++) {
        layer[k] = map<int, set<int>>();
        for (const auto &vertex : graph) {
            if (boundary.find(vertex.first) == boundary.end()) {
                map<int, int> count; // Count closed-by processors
                bool inlevel = false;
                
                for (const auto target: vertex.second) 
                    if (boundary.find(target) != boundary.end()) {
                        if (count.find(boundary[target]) != count.end())
                            count[boundary[target]] += 1;
                        else
                            count[boundary[target]] = 1;
                        inlevel = true;
                    }
                
                
                // If vertex is adjacent to the boundary set, update layer and boundary.
                if (inlevel) {
                    int maxdeg = 0;
                    int nearest;
                    for (const auto &v : count) 
                        if (v.second > maxdeg) {
                            maxdeg = v.second;
                            nearest = v.first;
                        }
                    
                    
                    if (layer[k].find(nearest) == layer[k].end())
                        layer[k][nearest] = set<int>();
                    
                    layer[k][nearest].insert(vertex.first);
                    boundary[vertex.first] = nearest;
                }
            }
        }
    }
    
    for (const auto &k : layer)
        for (const auto &p : k.second)
            (*bd_size)[p.first] += (int) p.second.size();
}

// Update boundary
void updateBD(map<int, int> &boundary, map<int, int> &position,
        map<int, set<int>> &graph, int rank, int size) {
    
    boundary = map<int, int>();
    if (rank == 0) {
        for (const auto &vertex : graph)
            for (int target : vertex.second)
                if (graph.find(target) == graph.end())
                    boundary[target] = position[target];
        
        for (int i=1; i<size; i++) {
            int msg_size;
            MPI_Status status;
            MPI_Probe(i, 401, MPI_COMM_WORLD, &status);
            MPI_Get_count(&status, MPI_INT, &msg_size);
            vector<int> row(msg_size, 0);
            MPI_Recv(&row[0], msg_size, MPI_INT, i, 401, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
            
            for (int j=0; j<msg_size; j++)
                row[j] = position[row[j]];
            MPI_Send(&row[0], row.size(), MPI_INT, i, 402, MPI_COMM_WORLD);
            
        }
        
    } else {
       vector<int> pos;
       for (const auto &vertex : graph)
       for (int target : vertex.second)
       if (graph.find(target) == graph.end())
           pos.push_back(target);
       
       vector<int> nodes = pos;
       
       MPI_Send(&pos[0], pos.size(), MPI_INT, 0, 401, MPI_COMM_WORLD);
       MPI_Recv(&pos[0], pos.size(), MPI_INT, 0, 402, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
       for (int j=0; j<pos.size(); j++)
           boundary[nodes[j]] = pos[j];
        
    }
    
}

// Call lp_solve to solve a LP problem; get the numbers of nodes that should be moved in L
void rebalancing(map<int, int> &position,
        vector<int> &alpha,
        vector<vector<int>> &L, int size) {
    
    lprec *lp;
    int *colno = NULL;
    int idx = 0;
    int Ncol = size * size - size;
    int ret = 0;
    REAL *row = NULL;
    
    int mu = position.size() / size;
    vector<int> B(size, mu);
    for (const auto &pos : position)
        B[pos.second]--;
    
    
    lp = make_lp(0, Ncol);
    
    if (lp == NULL)
        ret = 1;
    
    if (ret == 0) {
        for (int i=0; i<size; i++)
            for (int j=0; j<size; j++)
                if (i != j) {
                    char varname[16];
                    string s = "p" + to_string(i) + "p" + to_string(j);
                    strncpy(varname, s.c_str(), sizeof(varname));
                    varname[sizeof(varname) - 1] = 0;
                    set_col_name(lp, ++idx, varname);
                }
            
        
        colno = (int *) malloc(Ncol * sizeof(*colno));
        row = (REAL *) malloc(Ncol * sizeof(*row));
        if ((colno == NULL) || (row == NULL))
            ret = 2;
    }
    
    
    // Numbers of nodes constraint
    if (ret == 0) {
        set_add_rowmode(lp, TRUE);
        for (int i=0; i<Ncol; i++) {
            colno[0] = i + 1;
            row[0] = 1;
            
            if (!add_constraintex(lp, 1, row, colno, GE, 0))
                ret = 3;
            
            if (!add_constraintex(lp, 1, row, colno, LE, alpha[i]))
                ret = 3;
            
        }
    }
    
    // Flow constraint
    if (ret == 0) {
        set_add_rowmode(lp, TRUE);
        
        for (int i=0; i<size; i++) {
            idx = 0;
            for (int j=1; j<size; j++) {
                colno[idx] = (size - 1) * i + j;
                row[idx++] = -1;
                colno[idx] = i + (size - 1) * j - ((j > i) ? -1 : size - 1);
                row[idx++] = 1;
            }
            
            if (!add_constraintex(lp, idx, row, colno, (B[i] < 0) ? GE : LE, B[i]))
                ret = 3;
            if (!add_constraintex(lp, idx, row, colno, (B[i] < 0) ? LE : GE, (int) (B[i] * 0.8)))
                ret = 3;
            
        }
    }
    
    // Objective function
    if (ret == 0) {
        set_add_rowmode(lp, FALSE);
        for (int i=0; i<Ncol; i++) {
            colno[i] = i + 1;
            row[i] = -1;
        }
        
        if (!set_obj_fnex(lp, Ncol, row, colno))
            ret = 4;
    }
    
    if (ret == 0) {
        set_maxim(lp);
        set_verbose(lp, IMPORTANT);
        //set_verbose(lp, SEVERE);
        //write_LP(lp, stdout);
        ret = (solve(lp) == OPTIMAL) ? 0 : 5;
    }
    
    if (ret == 0) {
        get_variables(lp, row);
        int r = 0; int c = 1;
        for (int i=0; i<Ncol; i++) {
            L[r][c] = row[i];
            c++;
            if (c == size) {
                c = 0;
                r++;
            }
            if (c == r) c++;
        }
    }
    
    if (verbose) {
        if (ret == 0)
            cout << "Found a solution for the LP.\n";
        else
            cout << "Not found a solution for the LP.\n";
    }
    
    if (row != NULL)
        free(row);
    if (colno != NULL)
        free(colno);
    if (lp != NULL)
        delete_lp(lp);
    
}



void sampling(map<int, set<int>> &graph, map<int, int> &state, map<int, int> val) {
    for (const auto &vertex : graph) {
        int v = vertex.first;
        for (int i : vertex.second) {
            if (val.find(i) != val.end())
                state[v] += val[v];
            if (state.find(i) != state.end())
                state[v] += state[i];
        }
        poisson_distribution<int> rpois(state[v] * 1.0 / (graph[v].size() + 1));
        state[v] = min(rpois(gen), 1000);
    }
}

void receiveSampler(map<int, int> &state,
        map<int, set<int>> &graph, int i,
        map<int, int> &val) {
    
    int msg_size;
    MPI_Status status;
    MPI_Probe(i, 501, MPI_COMM_WORLD, &status);
    MPI_Get_count(&status, MPI_INT, &msg_size);
    vector<int> row(msg_size, 0);
    MPI_Recv(&row[0], msg_size, MPI_INT, i, 501, MPI_COMM_WORLD, &status);
    for (int j=0; j<msg_size; j+=2)
        val[row[j]] = row[j+1];
    
}

void sampler(map<int, set<int>> &graph,
        map<int, int> &boundary,
        map<int, int> &state,
        int rank, int size) {
    
    vector<vector<int>> msg(size, vector<int>());
    for (const auto &vertex : graph) {
        for (int target : vertex.second) {
            if (boundary.find(target) != boundary.end()) {
                int p = boundary[target];
                msg[p].push_back(target);
                msg[p].push_back(state[vertex.first]);
            }
        }
    }
    
    for (const auto &vertex : graph)
    for (int target : vertex.second)
    if (boundary.find(target) == boundary.end())
    state[vertex.first] += state[target];
    
    for (int k=0; k<100; k++) {
        map<int, int> val;
        for (int i=0; i<size; i++) {
            if (i < rank) {
                receiveSampler(state, graph, i, val);
                MPI_Send(&msg[i][0], msg[i].size(), MPI_INT, i, 501, MPI_COMM_WORLD);
                } else if (i > rank) {
                MPI_Send(&msg[i][0], msg[i].size(), MPI_INT, i, 501, MPI_COMM_WORLD);
                receiveSampler(state, graph, i, val);
            }
            
        }
        sampling(graph, state, val);
    }
    
}

void sendNodes(map<int,int> &position,
map<int, set<int>> &graph,
map<int, map<int, set<int>>> &layer,
map<int, int> &state,
vector<int> row,
int i, int rank) {
    
    int count = row[i];
    int level = 0;
    vector<int> cand;
    while (count > 0) {
        if (layer.find(level) == layer.end()) break;
        for (int j : layer[level][i]) {
            cand.push_back(j);
            count--;
            if (count == 0) break;
        }
        level++;
    }
    
    vector<int> msg;
    vector<int> curr;
    
    for (int j : cand) {
        msg.push_back(j);
        curr.push_back(state[j]);
        for (int target : graph[j])
        msg.push_back(target);
        msg.push_back(-1);
        state.erase(j);
        graph.erase(j);
        if (rank == 0)
        position[j] = i;
    }
    
    MPI_Send(&curr[0], curr.size(), MPI_INT, i, 202, MPI_COMM_WORLD);
    MPI_Send(&msg[0], msg.size(), MPI_INT, i, 201, MPI_COMM_WORLD);
    cand.push_back(i);
    
    if (rank > 0)
    MPI_Send(&cand[0], cand.size(), MPI_INT, 0, 203, MPI_COMM_WORLD);
    
}

void receiveNodes(map<int, int> &position,
map<int, set<int>> &graph,
map<int, int> &state,
int i, int rank) {
    
    int msg_size = 0;
    MPI_Status status;
    MPI_Probe(i, 201, MPI_COMM_WORLD, &status);
    MPI_Get_count(&status, MPI_INT, &msg_size);
    vector<int> row(msg_size, 0);
    MPI_Recv(&row[0], msg_size, MPI_INT, i, 201, MPI_COMM_WORLD, &status);
    
    MPI_Probe(i, 202, MPI_COMM_WORLD, &status);
    MPI_Get_count(&status, MPI_INT, &msg_size);
    vector<int> curr(msg_size, 0);
    MPI_Recv(&curr[0], msg_size, MPI_INT, i, 202, MPI_COMM_WORLD, &status);
    
    addGraph(row, &curr, graph, state, rank);
    if (rank == 0)
    for (int j=1; j<row.size(); j++) {
        if (row[j] < 0)
        j+=2;
        if (j < row.size())
        position[row[j]] = 0;
    }
}


void moveNodes(int rank, int size, vector<int> row,
map<int, set<int>> &graph,
map<int, map<int, set<int>>> &layer,
map<int, int> &state,
map<int, int> &position) {
    
    for (int i=0; i<size; i++)
    if (i > rank) {
        sendNodes(position, graph, layer, state, row, i, rank);
        receiveNodes(position, graph, state, i, rank);
        } else if (i < rank) {
        receiveNodes(position, graph, state, i, rank);
        sendNodes(position, graph, layer, state, row, i, rank);
    }
    
}


int main(int argc, char** argv) {
    // Varibles for MPI
    int rank, size;
    MPI_Status status;
    MPI_Comm comm;
    double time0=0, time11=0, time12=0;
    double time21=0, time22=0, time23=0, time24=0;
    double time3=0, time4=0;
    if (use_label) timeit = false;
    
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    int level = 5; // Tune the number of levels that vertices can move
    
    map<int, set<int>> graph; // graph[source] = map<target, weight>
    map<int, int> boundary; // <vertex, processor>
    map<int, int> position; // <vertex, processor>
    map<int, map<int, set<int>>> layer; // layer[level] = map<nearest_processor, set<vertex>>
    
    // Computation task
    map<int, int> state;
    
    if (rank == 0) {
        cout << "Started with " << size << " processors.\n";
        time0 = MPI_Wtime();
    }
    
    cout << "Processor " << rank << " is running.\n";
    
    // Initial partition, loaded from file
    // Graph.txt:
    // source partition target_1 target_2 ...
    readGraph("data/graph0.txt", rank, size, graph, position);
    updateBD(boundary, position, graph, rank, size);
    
    for (const auto &vertex : graph)
    state[vertex.first] = rpois0(gen);
    
    MPI_Barrier(MPI_COMM_WORLD);
    
    // Get new vertices
    for (int r=1; r<11; r++) {
        if (rank == 0) {
            if (timeit) time11 -= MPI_Wtime();

            readIncremental("data/graph" + to_string(r) + ".txt",
                          graph, boundary, position, state, size);

            if (verbose) cout << "=========== T = " << r << "\n";
            if (timeit) time11 += MPI_Wtime();
            if (timeit) time12 -= MPI_Wtime();

            } else {
                // Receive incremental graph
                int msg_size;
                MPI_Probe(0, 101, MPI_COMM_WORLD, &status);
                MPI_Get_count(&status, MPI_INT, &msg_size);
                if (msg_size > 0) {
                    vector<int> row(msg_size, 0);
                    MPI_Recv(&row[0], msg_size, MPI_INT, 0, 101, MPI_COMM_WORLD, &status);
                    addGraph(row, NULL, graph, state, rank);
                }
            
                MPI_Probe(0, 102, MPI_COMM_WORLD, &status);
                MPI_Get_count(&status, MPI_INT, &msg_size);
                if (msg_size > 0) {
                    vector<int> row(msg_size, 0);
                    MPI_Recv(&row[0], msg_size, MPI_INT, 0, 102, MPI_COMM_WORLD, &status);
                    addEdges(row, graph);
            }
            
        }
       
        
        if (timeit) {
            MPI_Barrier(MPI_COMM_WORLD);
            if (rank == 0) time12 += MPI_Wtime();
        }
        
        
        // Processor 0 finds the optimal allocation.
        // All processors send nodes to each other
        // and then update graph
        if (!use_label) {
            
            if (timeit && rank == 0) time21 -= MPI_Wtime();
            
            vector<int> alpha;
            vector<vector<int>> L(size, vector<int>(size, 0));
            
            // Update layer
            updateBD(boundary, position, graph, rank, size);
            
            vector<int> bd_size(size, 0);
            updateLayer(layer, graph, boundary, &bd_size, rank, level);
            
            
            if (rank > 0)
                MPI_Send(&bd_size[0], size, MPI_INT, 0, 301, MPI_COMM_WORLD);
            else {
                for (int i=0; i<size; i++)
                    if (i != rank)
                        alpha.push_back(bd_size[i]);
                for (int s=1; s<size; s++) {
                    MPI_Recv(&bd_size[0], size, MPI_INT, s, 301, MPI_COMM_WORLD, &status);
                    for (int i=0; i<size; i++)
                        if (i != s)
                            alpha.push_back(bd_size[i]);
                }
                if (timeit) time21 += MPI_Wtime();
            }
            // Processor 0 finds optimal transition.
            vector<int> row(size, 0);
            if (rank == 0) {
                if (timeit) time22 -= MPI_Wtime();
                    rebalancing(position, alpha, L, size);
                for (int i=1; i<size; i++)
                    MPI_Send(&L[i][0], size, MPI_INT, i, 302, MPI_COMM_WORLD);
                for (int i=0; i<size; i++)
                    row[i] = L[0][i];
                if (timeit) time22 += MPI_Wtime();
            } else
                MPI_Recv(&row[0], size, MPI_INT, 0, 302, MPI_COMM_WORLD, &status);
            
            // Sending nodes
            if (timeit) {
                MPI_Barrier(MPI_COMM_WORLD);
                if (rank == 0)
                    time23 -= MPI_Wtime();
            }
            
            MPI_Barrier(MPI_COMM_WORLD);
            
            if (timeit) time23 += MPI_Wtime();
            if (timeit) time24 -= MPI_Wtime();
            // All processors add received nodes.
            moveNodes(rank, size, row, graph, layer, state, position);
            
            if (timeit) {
                MPI_Barrier(MPI_COMM_WORLD);
                if (rank == 0) {
                    time24 += MPI_Wtime();
                    time3 -= MPI_Wtime();
                }
            }
            
            
            // Processor 0 updates positions.
            if (rank == 0) {
                for (int i=1; i<size; i++) {
                    int msg_size;
                    MPI_Probe(i, 203, MPI_COMM_WORLD, &status);
                    MPI_Get_count(&status, MPI_INT, &msg_size);
                    vector<int> row(msg_size, 0);
                    MPI_Recv(&row[0], msg_size, MPI_INT, i, 203, MPI_COMM_WORLD, &status);
                    int idx = row[msg_size - 1];
                    for (int j=0; j<msg_size-1; j++)
                        position[row[j]] = idx;
                }
                if (timeit) time3 += MPI_Wtime();
            }
            
        }
        
        if (timeit) {
            MPI_Barrier(MPI_COMM_WORLD);
            if (rank == 0)
                time4 -= MPI_Wtime();
        }
        
        updateBD(boundary, position, graph, rank, size);
        if (rank == 0)
            output("position" + to_string(r) + ".txt", position);
        
        sampler(graph, boundary, state, rank, size);
        
        if (timeit) {
            MPI_Barrier(MPI_COMM_WORLD);
            if (rank == 0)
                time4 += MPI_Wtime();
        }
        
        if (timeit && rank == 0) {
            cout << "Time: Step 1.1 " << time11 << "\n";
            cout << "Time: Step 1.2 " << time12 << "\n";
            cout << "Time: Step 2.1 " << time21 << "\n";
            cout << "Time: Step 2.2 " << time22 << "\n";
            cout << "Time: Step 2.3 " << time23 << "\n";
            cout << "Time: Step 2.4 " << time24 << "\n";
            cout << "Time: Step 3 " << time3 << "\n";
            cout << "Time: Step 4 " << time4 << "\n";
        }
        if (rank == 0)
            cout << "Total time " << MPI_Wtime() - time0 << "\n";
        
    }
    
    
    MPI_Barrier(MPI_COMM_WORLD);
    cout << "Done.\n";
    if (rank == 0) {
        cout << "Total time " << MPI_Wtime() - time0 << "\n";
        if (timeit) {
            cout << "Time: Step 1.1 " << time11 << "\n";
            cout << "Time: Step 1.2 " << time12 << "\n";
            cout << "Time: Step 2.1 " << time21 << "\n";
            cout << "Time: Step 2.2 " << time22 << "\n";
            cout << "Time: Step 2.3 " << time23 << "\n";
            cout << "Time: Step 2.4 " << time24 << "\n";
            cout << "Time: Step 3 " << time3 << "\n";
            cout << "Time: Step 4 " << time4 << "\n";
        }
    }
    
    MPI_Finalize();
    
    return 0;
}
